Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization
基于预条件空间的舍入：重新设计4-bit AdamW优化器状态量化
🧭 4-bit AdamW 优化器状态量化
University of California, Berkeley · Nubank
提出在预条件空间做随机舍入的4-bit AdamW状态量化方法，显著缩小与32-bit训练的验证损失差距
📌 工作参考价值（4 点）
- 若已经在用 TorchAO 4-bit AdamW 压缩大模型优化器状态，可直接替换第二动量的量化方式：不再做 state-space rounding，而是在 preconditioner space 计算随机舍入概率，改动集中在量化器实现层，训练流程无需重构。
- 第二动量码书可两条路线选择：ZIP-SR 保留零点适合大多数场景；ZE-EDEN 排除零点并 rescale 量化块，可缓解正量化底带来的预条件失真，对推荐/搜索模型中大量稀疏或接近零的梯度更有针对性。
- 第一动量用 4-bit NF4 已经够稳，但输出层（如推荐系统中的 item embedding 或 LLM 的 LM head）对量化更敏感，最后 10% 训练步对输出层第一动量做针对性随机舍入，能进一步降低验证损失，适合在微调生成式推荐模型时采用。
- 对需要大规模预训练或全参数 SFT 的电商/Agent 场景，该方法可降低优化器显存占用，同时把 4-bit 训练与 32-bit 训练之间的验证损失差距最多缩小 70%，适合作为训练成本优化的默认配置。